module Compiler where

import Control.Monad
import Control.Monad.State
import Data.List (intercalate)
import qualified Data.Map as Map
import Data.Map (Map, empty, insert, fromList)
import AbsJavalette


type VarTable = Map Ident Int
-- vars=variables, limitStackCounter=maxVar, stackCounter, labelCounter, tempCode, code , classname
data Env = Env {
	vars :: [Map Ident Int],
	localCounter :: Int,
	maxVar :: Int,
	stackCounter :: Int,
	maxStack :: Int,
	labelCounter :: Int,
	tempCode :: String,
	code :: [String],
	classname :: String
}

type EnvM = State Env

--initialzies an emtpy state
emtpyState :: String -> Env
emtpyState classname = Env [] 0 0 0 0 0 "" [] classname

compileProgram :: Program -> String -> String
compileProgram (Program p) classname = do
	let code = evalState (compileFunctions p >> getCode) (emtpyState classname)
	classDefinition classname code

compileFunctions :: [TopDef] -> EnvM ()
compileFunctions [] = return ()
compileFunctions (f:fs) = do
	compileFunction f
	compileFunctions fs
	return ()

compileFunction :: TopDef -> EnvM ()
compileFunction f@(FnDef t (Ident i) args (Block b)) = do
	newBlock
	compileFunctionArguments args
	compileStms b
	doFinalCode f
	return ()

compileFunctionArguments :: [Arg] -> EnvM ()
compileFunctionArguments [] = return ()
compileFunctionArguments ((Arg t i):as) = do
	addVar i t
	compileFunctionArguments as

doFinalCode:: TopDef -> EnvM ()
doFinalCode (FnDef typ (Ident identifier) args (Block b)) = do
	hackRet typ
	let p = intercalate "" (map (\(Arg t _) -> jvmType t) args)
	Env v l mv s msc lc tc c cln <- get
	let newFinal = c ++
		[".method public static " ++ identifier ++ "(" ++ p ++ ")" ++ (jvmType typ) ++ "\n" ++
		" .limit locals " ++ (show l) ++ "\n" ++ 
		" .limit stack " ++ (show msc) ++ "\n" ++
		tc ++ 
		".end method\n\n"]
	put (Env [] 0 0 0 0 lc "" newFinal cln)
	return ()

classDefinition :: String -> String -> String
classDefinition classname code =
	".class public " ++ classname ++ "\n" ++
	".super java/lang/Object\n\n" ++
	".method public static main([Ljava/lang/String;)V\n" ++
	" .limit locals 1\n" ++
	" invokestatic " ++ classname ++ "/main()I\n" ++
	" pop\n" ++
	" return\n" ++
	".end method\n\n" ++ code

compileStms :: [Stmt] -> EnvM ()
compileStms [] = return ()
compileStms (s:ss) = do
	compileStm s
	compileStms ss
	return ()

compileArgs :: [Expr] -> EnvM ()
compileArgs [] = return ()
compileArgs (e:es) = do
	compileExpr e
	compileArgs es


compileItems :: Type -> [Item] -> EnvM ()
compileItems t [] = return ()
compileItems t (i:is) = do
	case i of
		NoInit identifier -> do
			addVar identifier t
			vLoc <- lookupVar identifier
			case t of
				Int -> do
					addTemproraryCode ("  bipush 0\n")
					incStackCounter 1
					addTemproraryCode ("  istore " ++ show vLoc ++ "\n")
					decStackCounter 1
				Bool -> do
					addTemproraryCode ("  bipush 0\n")
					incStackCounter 1
					addTemproraryCode ("  istore " ++ show vLoc ++ "\n")
					decStackCounter 1
				Doub -> do
					addTemproraryCode ("  ldc2_w 0.0\n")
					incStackCounter 2
					addTemproraryCode ("  dstore " ++ show vLoc ++ "\n")
					decStackCounter 2
			compileItems t is
		Init identifier expr -> do
			compileExpr expr
			addVar identifier t
			vLoc <- lookupVar identifier
			case t of
				Int -> do
					addTemproraryCode ("  istore " ++ show vLoc ++ "\n")
					decStackCounter 1
				Bool -> do
					addTemproraryCode ("  istore " ++ show vLoc ++ "\n")
					decStackCounter 1
				Doub -> do
					addTemproraryCode ("  dstore " ++ show vLoc ++ "\n")
					decStackCounter 2
			compileItems t is


compileStm :: Stmt -> EnvM ()
compileStm stmt = case stmt of
	Empty -> return ()
	BStmt (Block stmts) -> do
		newBlock
		compileStms stmts
		exitBlock
	Decl t items -> do
		compileItems t items
		return ()
	Ass identifier e -> do
		sl <- lookupVar identifier --get locatio of var
		compileExpr e
		case typeOf e of
			Int -> do
				addTemproraryCode ("  istore " ++ show sl ++ "\n")
				decStackCounter 1
			Bool -> do
				addTemproraryCode ("  istore " ++ show sl ++ "\n")
				decStackCounter 1
			Doub -> do
				addTemproraryCode ("  dstore " ++ show sl ++ "\n")
				decStackCounter 2
	Incr identifier -> do
		varL <- lookupVar identifier
		addTemproraryCode ("  iinc " ++ show varL ++ " 1\n")
		return()
	Decr identifier -> do
		varL <- lookupVar identifier
		addTemproraryCode ("  iinc " ++ show varL ++ " -1\n")
		return()
	Ret e -> do
		compileExpr e
		case typeOf e of
			Int -> do
				addTemproraryCode (returnType Int)
				decStackCounter 1
			Bool -> do
				addTemproraryCode (returnType Bool)
				decStackCounter 1
			Doub -> do
				addTemproraryCode (returnType Doub)
				decStackCounter 2
		return ()
	VRet -> do
		addTemproraryCode ("  return\n")
		return()
	Cond e s -> do
		endL <- getNewLabel
		compileExpr e
		addTemproraryCode ("  ifeq " ++ endL ++ "\n")
		decStackCounter 1
		compileStm s
		addTemproraryCode  ("  " ++ (endL ++ ":\n"))
		return ()
	CondElse e s1 s2 -> do
		falseLabel <- getNewLabel
		trueLabel <- getNewLabel
		compileExpr e
		addTemproraryCode ("  ifeq " ++ falseLabel ++ "\n")
		decStackCounter 1 -- ifeq consumes top of stack
		compileStm s1
		addTemproraryCode ("  goto " ++ trueLabel ++ "\n")
		addTemproraryCode ("  " ++ falseLabel ++ ":\n")
		compileStm s2
		addTemproraryCode ("  " ++ trueLabel ++ ":\n")
		return ()
	While e s -> do
		testLabel <- getNewLabel
		endLabel <- getNewLabel
		addTemproraryCode ("  " ++ testLabel ++ ":\n")
		compileExpr e
		addTemproraryCode ("  ifeq " ++ endLabel ++ "\n")
		decStackCounter 1 -- ifeq consuming top of stack
		compileStm s
		addTemproraryCode ("  goto " ++ testLabel ++ "\n")
		addTemproraryCode ("  " ++ endLabel ++ ":\n")
		return ()
	SExp e -> do
		compileExpr e
		return ()
	
compileExpr :: Expr -> EnvM ()
compileExpr expr = case expr of
	ETyped _ (EApp (Ident "printString") args) -> do
		let outString =  typedEStringToString (head args)
		addTemproraryCode ("  ldc " ++ show outString ++ "\n")
		addTemproraryCode ("  invokestatic Runtime/printString(Ljava/lang/String;)V\n")
		incStackCounter 1 --Loading the string lit onto the stack increases the stack size
		decStackCounter 1 --printString consumes the top of the stack, we need to increase and then decrease the stack count so that the max stack count may be bumped up if needed
		return ()
	ETyped _ (EApp (Ident "printInt") args) -> do
		let arg = head args
		compileExpr arg
		addTemproraryCode ("  invokestatic Runtime/printInt(I)V\n")
		decStackCounter 1
		return ()
	ETyped _ (EApp (Ident "printDouble") args) -> do
		let arg = head args
		compileExpr arg
		addTemproraryCode ("  invokestatic Runtime/printDouble(D)V\n")
		decStackCounter 2
		return ()
	ETyped Int (EApp (Ident "readInt") _) -> do
		addTemproraryCode ("  invokestatic Runtime/readInt()I\n")
		incStackCounter 1
		return ()
	ETyped Doub (EApp (Ident "readDouble") _) -> do
		addTemproraryCode ("  invokestatic Runtime/readDouble()D\n")
		incStackCounter 2
		return ()
	ETyped t (EApp (Ident identifier) args) -> do
		compileArgs args
		cn <- getClassName
		addTemproraryCode ("  invokestatic " ++ cn ++ "/" ++ identifier ++ "(" ++ (argTypes args) ++ ")" ++ (jvmType t) ++ "\n")
		case t of
			Int -> incStackCounter 1
			Bool -> incStackCounter 1
			Doub -> incStackCounter 2
			Void -> return ()
	ETyped t (EVar i) -> do
		sl <- lookupVar i -- find location of var
		case t of
			Int -> do 
				addTemproraryCode ("  iload " ++ show sl ++ "\n")
				incStackCounter 1
			Bool -> do 
				addTemproraryCode ("  iload " ++ show sl ++ "\n")
				incStackCounter 1
			Doub -> do
				addTemproraryCode ("  dload " ++ show sl ++ "\n")
				incStackCounter 2		
	ETyped t ELitTrue -> do
		addTemproraryCode ("  bipush " ++ "1\n")
		incStackCounter 1
		return ()
	ETyped t ELitFalse -> do
		addTemproraryCode ("  bipush " ++ "0\n")
		incStackCounter 1
		return ()
	ETyped t (ELitInt v) -> do
		addTemproraryCode ("  ldc " ++ show v ++ "\n") -- we could add optimizations checking for the size of the int and using bipush for 1 byte ints
		incStackCounter 1
		return ()
	ETyped t (ELitDoub d) -> do
		addTemproraryCode ("  ldc2_w " ++ show d ++ "\n")
		incStackCounter 2
		return ()
	ETyped t (EMul a op b) -> case op of
		Times -> do
			compileExpr a
			compileExpr b
			case t of
				Int -> do 
					addTemproraryCode("  imul\n")
					decStackCounter 1
				Doub -> do
					addTemproraryCode("  dmul\n")
					decStackCounter 2
			return ()
		Div -> do
			compileExpr a
			compileExpr b
			case t of
				Int -> do
					addTemproraryCode("  idiv\n")
					decStackCounter 1
				Doub -> do
					addTemproraryCode("  ddiv\n")
					decStackCounter 2
		Mod -> do
			compileExpr a
			compileExpr b
			case t of
				Int -> do --should we have a case for doubles as well? i think there is an instruction for it in jasmins instruction set
					addTemproraryCode("  irem\n")
					decStackCounter 1
	ETyped t (EAdd a op b) -> case op of
		Plus -> do 
			compileExpr a
			compileExpr b
			case t of
				Int -> do
					addTemproraryCode("  iadd\n")
					decStackCounter 1
				Doub -> do 
					addTemproraryCode("  dadd\n")
					decStackCounter 2
			return ()
		Minus -> do
			compileExpr a
			compileExpr b
			case t of
				Int -> do
					addTemproraryCode("  isub\n")
					decStackCounter 1
				Doub -> do
					addTemproraryCode("  dsub\n")
					decStackCounter 2
			return ()
	ETyped t (Neg e) -> do
		compileExpr e
		case t of
			Int -> do
				addTemproraryCode("  ineg\n")
			Doub -> do
				addTemproraryCode("  dneg\n")
		return ()
	ETyped t (Not e) -> do 
		tLabel <- getNewLabel
		addTemproraryCode ("  bipush 1\n")
		incStackCounter 1
		compileExpr e
		addTemproraryCode ("  ifeq " ++ tLabel ++ "\n")
		decStackCounter 1
		addTemproraryCode ("  pop\n")
		addTemproraryCode ("  bipush 0\n");
		addTemproraryCode ("  " ++ tLabel ++ ":\n")
		return ()
	ETyped t (ERel e1 op e2) -> do  --the results of theese should be an boolean on the stack (int 0 or 1)
		case typeOf e1 of
			Int -> do
				trueLabel <- getNewLabel
				addTemproraryCode ("  bipush 1\n")
				incStackCounter 1
				compileExpr e1
				compileExpr e2
				case op of
					LTH -> addTemproraryCode ("  if_icmplt " ++ trueLabel ++ "\n")
					LE -> addTemproraryCode ("  if_icmple " ++ trueLabel ++ "\n")
					GTH -> addTemproraryCode ("  if_icmpgt " ++ trueLabel ++ "\n")
					GE -> addTemproraryCode ("  if_icmple " ++ trueLabel ++ "\n")
					NE -> addTemproraryCode ("  if_icmpne " ++ trueLabel ++ "\n")
					EQU -> addTemproraryCode ("  if_icmpeq " ++ trueLabel ++ "\n")
				decStackCounter 2 -- i guess if_icmplt consumes the two top values of the stack
				addTemproraryCode ("  pop\n")
				addTemproraryCode ("  bipush 0\n")
				addTemproraryCode ("  " ++ trueLabel ++ ":\n")
				return ()
			Bool -> do
				trueLabel <- getNewLabel
				addTemproraryCode ("  bipush 1\n")
				incStackCounter 1
				compileExpr e1
				compileExpr e2
				case op of
					LTH -> addTemproraryCode ("  if_icmplt " ++ trueLabel ++ "\n")
					LE -> addTemproraryCode ("  if_icmple " ++ trueLabel ++ "\n")
					GTH -> addTemproraryCode ("  if_icmpgt " ++ trueLabel ++ "\n")
					GE -> addTemproraryCode ("  if_icmple " ++ trueLabel ++ "\n")
					NE -> addTemproraryCode ("  if_icmpne " ++ trueLabel ++ "\n")
					EQU -> addTemproraryCode ("  if_icmpeq " ++ trueLabel ++ "\n")
				decStackCounter 2 -- i guess if_icmplt consumes the two top values of the stack
				addTemproraryCode ("  pop\n")
				addTemproraryCode ("  bipush 0\n")
				addTemproraryCode ("  " ++ trueLabel ++ ":\n")
				return ()
			Doub -> do --TODO LE GE
				trueLabel <- getNewLabel
				addTemproraryCode ("  bipush 1\n")
				incStackCounter 1
				compileExpr e1
				compileExpr e2	
				addTemproraryCode ("  dcmpg\n") -- top of the stack is now -1, 0 or 1 
				decStackCounter 3 --two doubles gone and one int added
				case op of
					LTH -> do
						addTemproraryCode ("  bipush -1\n")
						addTemproraryCode ("  if_icmpeq " ++ trueLabel ++ "\n")
					GTH -> do
						addTemproraryCode ("  bipush 1\n")
						addTemproraryCode ("  if_icmpeq " ++ trueLabel ++ "\n")
					EQU -> do
						addTemproraryCode ("  bipush 0\n")
						addTemproraryCode ("  if_icmpeq " ++ trueLabel ++ "\n")
					LE -> do
						addTemproraryCode ("  bipush -1\n")
						addTemproraryCode ("  if_icmple " ++ trueLabel ++ "\n")
					GE -> do
						addTemproraryCode ("  bipush 1\n")
						addTemproraryCode ("  if_icmpge " ++ trueLabel ++ "\n")
				decStackCounter 1
				addTemproraryCode ("  pop\n")
				addTemproraryCode ("  bipush 0\n")
				addTemproraryCode ("  " ++ trueLabel ++ ":\n")
				return ()
	ETyped t (EAnd e1 e2) -> do
		eLabel <- getNewLabel
		addTemproraryCode ("  bipush 0\n") --The default return value
		incStackCounter 1 
		compileExpr e1 --the top of the stack now contains 0 if the expression is false else 1
		addTemproraryCode ("  ifeq " ++ eLabel ++ "\n") -- if the top of the stack is 0 (e1 false) then jump to the end, the test is finished
		decStackCounter 1 --ifeq consumes the top of the stack
		compileExpr e2
		addTemproraryCode ("  ifeq " ++ eLabel ++ "\n") -- if the top of the stack is 0 (e2 false) then jump to the end, the test is finished
		decStackCounter 1
		addTemproraryCode ("  pop\n")
		addTemproraryCode ("  bipush 1\n")
		addTemproraryCode ("  " ++ eLabel ++ ":\n")
		return ()
	ETyped t (EOr e1 e2) -> do
		eLabel <- getNewLabel
		addTemproraryCode ("  bipush 1\n") --assume true and change to false only if both expressions are false
		incStackCounter 1
		compileExpr e1
		addTemproraryCode ("  ifne " ++ eLabel ++ "\n")
		compileExpr e2
		addTemproraryCode ("  ifne " ++ eLabel ++ "\n")
		decStackCounter 1
		addTemproraryCode ("  pop\n") --if we get this far, then both e1 and e2 evaluated to false replace the default true on the stack with false
		addTemproraryCode ("  bipush  0\n")
		addTemproraryCode ("  " ++ eLabel ++ ":\n")


returnType :: Type -> String
returnType t = case t of
	Int -> "  ireturn\n"
	Bool -> "  ireturn\n"
	Doub -> "  dreturn\n"
	Void -> "  return\n"


--Hack to get the string needed for application of the printString function
typedEStringToString :: Expr -> String
typedEStringToString (ETyped String (EString s)) = s

argTypes :: [Expr] -> String
argTypes args = concatenation (map jvmType (map typeOf args))

jvmType :: Type -> String
jvmType Int = "I"
jvmType Doub = "D"
jvmType Void = "V"
jvmType Bool = "I"
jvmType String = "Ljava/lang/String;"

concatenation :: [String] -> String
concatenation [] = ""
concatenation (s:ss) = s ++ concatenation ss

sizeOf :: Type -> Int
sizeOf Int = 1
sizeOf Bool = 1
sizeOf Doub = 2

typeOf :: Expr -> Type
typeOf (ETyped t _) = t
--state helper functions

addVar :: Ident -> Type -> EnvM ()
addVar identifier typ = do
	Env (v:vs) l mv s msc lc tc c cln <- get
	let v' = (insert identifier l v)
	put (Env (v':vs) (l+(sizeOf typ)) mv s msc lc tc c cln)
	return ()

lookupVar :: Ident -> EnvM (Int)
lookupVar identifier = do
	Env v _ _ _ _ _ _ _ _ <- get
	return (findInVarStack v identifier)

findInVarStack :: [Map Ident Int] -> Ident -> Int
findInVarStack [] (Ident identifier) = error ("var not found " ++ identifier)
findInVarStack (t:rest) identifier = case (Map.lookup identifier t) of
	Just i -> i
	Nothing -> findInVarStack rest identifier

newBlock :: EnvM ()
newBlock = do
	Env v l mv s msc lc tc c cln <- get
	put (Env (empty:v) l mv s msc lc tc c cln)
	return ()

exitBlock :: EnvM ()
exitBlock = do
	Env (v:vs) l mv s msc lc tc c cln <- get
	put (Env vs l mv s msc lc tc c cln)


getNewLabel :: EnvM (String)
getNewLabel = do
	Env v l mv s msc lc tc c cln <- get
	put (Env v l mv s msc (lc+1) tc c cln)
	return ("label" ++ (show lc))

incStackCounter :: Int -> EnvM ()
incStackCounter i = do 
	Env v l mv s msc lc tc c cln <- get
	let s' = s+i
	if(s' > msc) then
		put (Env v l mv (s+i) s' lc tc c cln)
	else
		put (Env v l mv (s+i) msc lc tc c cln)
	return ()

decStackCounter :: Int -> EnvM ()
decStackCounter i = do 
	Env v l mv s msc lc tc c cln <- get
	put (Env v l mv (s-i) msc lc tc c cln)
	return ()

addTemproraryCode :: String -> EnvM ()
addTemproraryCode newCode = do 
	Env v l mv s msc lc tc c cln <- get
	put (Env v l mv s msc lc (tc++newCode) c cln)
	return ()

getClassName :: EnvM (String)
getClassName = do
	Env v l mv s msc lc tc c cln <- get
	return cln

getCode :: EnvM (String)
getCode = do
	Env v l mv s msc lc tc c cln <- get
	return (concatenation c)


hackRet :: Type -> EnvM () --this is bullshit, we know, but we have run out of time to do it properly
hackRet t = case t of
	Void -> do
		addTemproraryCode ("  return\n")
	Int -> do 
		addTemproraryCode ("  bipush 0\n")
		incStackCounter 1
		addTemproraryCode ("  ireturn\n")
		decStackCounter 1
	Bool -> do
		addTemproraryCode ("  bipush 0\n")
		incStackCounter 1
		addTemproraryCode ("  ireturn\n")
		decStackCounter 1
	Doub -> do
		addTemproraryCode ("  ldc2_w 0.0\n")
		incStackCounter 2
		addTemproraryCode ("  dreturn\n")
		decStackCounter 2
